from ultralytics import YOLO
import  cv2
import time
import mypose

model = YOLO('yolo11n-pose.pt')

cap = cv2.VideoCapture(0)
frameWidth = cap.get(cv2.CAP_PROP_FRAME_WIDTH)#cv2.CAP_PROP_FRAME_WIDTH,
frameHeight = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)#cv2.CAP_PROP_FRAME_HEIGHT,
fps = cap.get(cv2.CAP_PROP_FPS)#cv2.CAP_PROP_FPS,

out_path = 'output.mp4'
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
# 需要整数
out = cv2.VideoWriter(out_path, fourcc, fps, (int(frameWidth), int(frameHeight)))



while True:
    ret, frame = cap.read()
    if not ret:
        break
    results = model(frame)[0]
       
    # 创建PoseKeypoints实例
    if len(results) >= 2:
        pose1 = mypose.PoseKeypoints(results, frameWidth, frameHeight,person_id=0)
        pose2 = mypose.PoseKeypoints(results, frameWidth, frameHeight, person_id=1)

    # 调用PoseShow方法绘制骨架
        skl1 = pose1.PoseShow(frame,show_angles=True)
        skl2 = pose2.PoseShow(frame,  show_angles=True)

    
    # 比较两个人姿态的相似度
        mypose.compare_poses(frame, pose1, pose2)

        out.write(frame)
        cv2.imshow('frame', frame)
    else:
        print("图像中只有一个人，无法进行动作评分")

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
out.release()
cv2.destroyAllWindows()